# https://www.bilibili.com/video/BV17w411k7a9/?spm_id_from=333.1007.top_right_bar_window_custom_collection.content.click&vd_source=b664701fe1351134c3243c7903d11516
# 数据集 MINIST 里面都是一张一张的图片 
# 每张图是28*28像素(即784个格子)，灰度图，只有一个通道，每个像素点(格子)只有一个数值，
# 所以每张图(每个样本)，都有784个数字
# 每张图宏观显示都是手写的一个数字
# 所以，输出层为10个 ，要预测改2样本为数字0～9的哪一个
# 所以一个样本有784个特征点，所以全连接神经网络输入层数量为784 即x_1_1,x_1_2,...,x_784_1
import numpy as np #要用的MINIST数据集 是numpy格式
import torch
import gzip
import argparse
import  torch.nn as nn
import torch.nn.functional as F #激活函数 交叉熵损失函数都在这个里面
import os  # 添加os模块用于路径操作
#定义超参数
def parse_option():
    parser = argparse.ArgumentParser(description='MNIST MLP Training')
    # parser.add_argument('--epochs', type=int, help='the number of epoch', default=3) #epochs训练次数
    parser.add_argument("--save_path", type=str, help="the path of model saved",
                        default='models/MLP.pth')
    args = parser.parse_args()  # 也可直接使用 args, _ = parser.parse_known_args()
    return args


def load_mnist_images(filename):
    """加载MNIST图像文件"""
    with gzip.open(filename, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=16)
    data = data.reshape(-1, 28, 28)
    return data

def load_mnist_labels(filename):
    """加载MNIST标签文件"""
    with gzip.open(filename, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=8)
    return data

# 创建神经网络模型对应的类    
class MNIST(nn.Module):# 继承nn.Module
    def __init__(self): # 成员函数：初始化方法，可以知道当前定义的类 之后实例出来的对象(我们的模型)里面可以包含有哪些属性(成员变量)和方法(成员函数)
        super().__init__() #继承父类的东西
        self.fc1 = nn.Linear(784,1000) #类的属性 第一个全连接层 1000为第二层的神经元数量  用nn里面的Linear帮我们搭建(Linear帮忙完成了wx+b的事情)
        self.fc2 = nn.Linear(1000,500)
        self.fc3 = nn.Linear(500,10)
    def forward(self, x): #前向传播函数
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x) # 不用softmax，因为后面交叉熵损失函数包含了softmax
                        # 最后一层的x通过softmax 求出p， p再去和标签y 用交叉熵求出loss
        return x

def main(args):
        # 确保保存目录存在
    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
    # 加载训练集和测试集
    train_images = load_mnist_images('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/train-images-idx3-ubyte.gz')
    train_labels = load_mnist_labels('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/train-labels-idx1-ubyte.gz')
    test_images = load_mnist_images('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/t10k-images-idx3-ubyte.gz')
    test_labels = load_mnist_labels('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/t10k-labels-idx1-ubyte.gz')

    # 使用字典保存，避免形状不一致的问题
    mnist_data = {
        'train_x': train_images,
        'test_x': test_images,
        'train_y': train_labels,
        'test_y': test_labels
    }

    # 保存为npy文件
    np.save('mnist.npy', mnist_data)
    print("MNIST数据已保存为 mnist.npy")

    #加载数据集
    # 解包数据集 训练集的数据，测试集的数据，训练集的标签，测试集的标签
    mnist_data = np.load('mnist.npy', allow_pickle=True).item()
    train_x = mnist_data['train_x']
    test_x = mnist_data['test_x']
    train_y = mnist_data['train_y']
    test_y = mnist_data['test_y']
    print(f"训练集图像形状: {train_x.shape}")  # (60000, 28, 28) 28行28列
    print(f"训练集标签形状: {train_y.shape}")   # (60000,)
    print(f"测试集图像形状: {test_x.shape}")   # (10000, 28, 28)
    print(f"测试集标签形状: {test_y.shape}")    # (10000,)

    x_train = train_x.reshape(60000,784).astype(np.float32)   
    print(f"训练集图像形状: {x_train.shape}")  # 转换成向量 784列


    # np的数据类型转换成pytorch的tenser类型
    x_train = torch.tensor(x_train)
    y_train = torch.tensor(train_y)
    print(f"*****************************")  
    print(f"训练集图像: {x_train.shape}")  # 转换成向量 784列
    print(f"训练集标签: {y_train.shape}")  


    print(f"*******************************搭建全连接神经网络*********************************") 
    # import  torch.nn as nn
    # import torch.nn.functional as F #激活函数 交叉熵损失函数都在这个里面
    # # 创建神经网络模型对应的类    
    # class MNIST(nn.Module):# 继承nn.Module
    #     def __init__(self): # 成员函数：初始化方法，可以知道当前定义的类 之后实例出来的对象(我们的模型)里面可以包含有哪些属性(成员变量)和方法(成员函数)
    #         super().__init__() #继承父类的东西
    #         self.fc1 = nn.Linear(784,1000) #类的属性 第一个全连接层 1000为第二层的神经元数量  用nn里面的Linear帮我们搭建(Linear帮忙完成了wx+b的事情)
    #         self.fc2 = nn.Linear(1000,500)
    #         self.fc3 = nn.Linear(500,10)
    #     def forward(self, x): #前向传播函数
    #         x = F.relu(self.fc1(x))
    #         x = F.relu(self.fc2(x))
    #         x = self.fc3(x) # 不用softmax，因为后面交叉熵损失函数包含了softmax
    #                         # 最后一层的x通过softmax 求出p， p再去和标签y 用交叉熵求出loss
    #         return x

    # 实例化
    mnist = MNIST() #mnist这个对象 就是个神经网络模型了
    y = mnist(x_train) #求出了模型的输出


    print(f"*******************************求损失函数*********************************") 
    loss = F.cross_entropy(y, y_train)

    print(f"*******************************反向传播,求每个w的梯度值grad********************************") 
    loss.backward()
    lr = 0.001 #学习率


    print(f"*******************************更新参数w(权重)********************************") 
    with torch.no_grad(): # 上下文管理器
        for w in mnist.parameters(): # 找到每一个模型的参数w
            w -= lr * w.grad # 梯度下降

    print(f"*******************************训练 更新参数w(权重) 100次********************************") 
    for i in range(100): # 训练100次 参数更新100次
        y = mnist(x_train) #求出了模型的输出
        loss = F.cross_entropy(y, y_train)
        loss.backward()
        with torch.no_grad(): # 上下文管理器
            for w in mnist.parameters(): # 找到每一个模型的参数w
                w -= lr * w.grad # 梯度下降
        mnist.zero_grad() #该函数作用：梯度清零
        # if i%5==0: #每5次迭代 就进入
        print(i,loss) # 第一个loss如果很小 说明初始化的w 初始化的很合适
        torch.save(mnist.state_dict(), args.save_path) #保存形式 将参数进行保存
    
#调用  入口函数
if __name__ == '__main__':
    args_ = parse_option()
    main(args_)